# ====================================================================================
# 文件: preprocessing/base_pre_parser.py
# 描述: [V4 修复版]
#      (修复) get_pre_config 现在接受一个可选的 argv_list 参数。
#      这允许 step1.py 先解析 --vllm-host，再将剩余参数
#      传递给此解析器，避免 "未知参数" 警告。
# ====================================================================================

import argparse
import os
import sys
from typing import Any, Dict

import torch
import yaml


def _dict_to_namespace(d):
    """
    递归地将字典转换为 argparse.Namespace
    """
    if not isinstance(d, dict):
        return d

    namespace = argparse.Namespace()
    for key, value in d.items():
        if isinstance(value, dict):
            setattr(namespace, key, _dict_to_namespace(value))
        elif isinstance(value, list):
            setattr(namespace, key, [
                _dict_to_namespace(item) if isinstance(item, dict) else item
                for item in value
            ])
        else:
            setattr(namespace, key, value)
    return namespace


def _flatten_train_config(config):
    if not hasattr(config, "train_config"):
        return

    train_cfg = config.train_config
    for key, value in vars(train_cfg).items():
        if key == "fusion_weights":
            continue
        setattr(config, key, value)

    fusion_cfg = getattr(train_cfg, "fusion_weights", None)
    if fusion_cfg:
        for key, value in vars(fusion_cfg).items():
            setattr(config, key, value)

    if hasattr(train_cfg, "user_recon"):
        setattr(config, "user_recon", _namespace_to_dict(train_cfg.user_recon))


def _namespace_to_dict(ns: Any) -> Dict[str, Any]:
    if not hasattr(ns, "__dict__"):
        return ns
    result = {}
    for key, value in vars(ns).items():
        if hasattr(value, "__dict__"):
            result[key] = _namespace_to_dict(value)
        else:
            result[key] = value
    return result


def get_pre_config(argv_list=None):  # [V14 修复 2/3] 接受参数列表
    """
    解析命令行参数 (仅 --config) 并加载 YAML 配置文件。
    """
    parser = argparse.ArgumentParser(description="SDKR Preprocessing Script")

    parser.add_argument('--config',
                        type=str,
                        required=True,
                        help='YAML 配置文件路径 (例如: configs/amazon-book.yaml)')

    # [V14 修复 2/3]
    # 如果 argv_list 未提供 (step2, step3)，则解析 sys.argv
    # 如果 argv_list 已提供 (step1)，则只解析该列表
    args_list_to_parse = argv_list if argv_list is not None else sys.argv[1:]

    args, unknown = parser.parse_known_args(args_list_to_parse)

    if unknown:
        # 这个警告现在只会在 step2/step3 收到未知参数时触发
        print(f"[PreParser] 警告: 忽略了以下未知参数: {unknown}")

    config_file_path = args.config

    if not os.path.exists(config_file_path):
        alt_path = os.path.join(os.path.dirname(__file__), '..', config_file_path)
        if os.path.exists(alt_path):
            config_file_path = alt_path
        else:
            raise FileNotFoundError(f"配置文件未找到: {args.config} (或 {alt_path})")

    print(f"[PreParser] 加载配置文件: {config_file_path}")

    with open(config_file_path, 'r') as f:
        config_dict = yaml.safe_load(f)

    config = _dict_to_namespace(config_dict)

    config.data_config = config.data
    config.preproc_config = config.preprocessing
    config.train_config = config.training

    config.PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(config_file_path), '..'))
    config.DATA_ROOT = os.path.join(config.PROJECT_ROOT, config.data_config.path)
    config.DATA_DIR = os.path.join(config.DATA_ROOT, config.data_config.dataset)

    print(f"[PreParser] 项目根目录: {config.PROJECT_ROOT}")
    print(f"[PreParser] 数据目录: {config.DATA_DIR}")

    _flatten_train_config(config)

    device_str = f"cuda:{config.gpu_id}" if torch.cuda.is_available() else "cpu"
    setattr(config, "device", torch.device(device_str))

    if hasattr(config, "top_k"):
        topk = getattr(config, "top_k")
        if isinstance(topk, int):
            topk = [topk]
        elif isinstance(topk, tuple):
            topk = list(topk)
        setattr(config, "top_k", topk)

    return config